import numpy as np
from pyDOE2 import lhs
import pandas as pd
from concurrent.futures import ProcessPoolExecutor
import modelABCDEF
import time


def train_and_evaluate(
    i, cumulative_samples_list_train, cumulative_samples_list_eval, n_eval_samples
):
    # 训练模型
    model1, model2 = modelABCDEF.train_model(cumulative_samples_list_train[i])
    results = []
    # 使用模型评价新采样的数据
    for j in range(n_eval_samples):
        result = modelABCDEF.use_model(cumulative_samples_list_eval[j], model1, model2)
        results.append(result)
    return results


if __name__ == "__main__":
    time_start = time.time()
    np.random.seed(42)
    # 定义年份范围
    years = list(range(2024, 2031))  # 从2024到2030年

    # 定义采样的样本数和变量数
    n_train_samples = 2  # 用于训练的样本数量
    n_eval_samples = 2  # 用于评价的样本数量
    n_variables = 4  # 我们有四个变量需要考虑

    # 第一次采样 - 用于训练模型
    lhs_samples_train = lhs(n_variables, samples=n_train_samples)

    # 定义每个变量的范围
    growth_rate_wheat_corn_train = 0.05 + lhs_samples_train[:, 0] * 0.05
    sales_variation_other_crops_train = -0.05 + lhs_samples_train[:, 1] * 0.10
    yield_variation_crops_train = -0.10 + lhs_samples_train[:, 2] * 0.20
    price_decrease_mushroom_train = 0.01 + lhs_samples_train[:, 3] * 0.04
    # 生成包含字典的列表，用于模型训练
    cumulative_samples_list_train = []
    for i in range(n_train_samples):
        year_data = {}
        cumulative_growth_wheat_corn = 1.0
        cumulative_sales_variation = 1.0
        cumulative_yield_variation = 1.0
        cumulative_price_decrease = 1.0
        yangdu = 1.0
        cost = 1.0
        vegetable_price = 1.0
        for year in years:
            cumulative_growth_wheat_corn *= 1 + growth_rate_wheat_corn_train[i]
            cumulative_sales_variation *= 1 + sales_variation_other_crops_train[i]
            cumulative_yield_variation *= 1 + yield_variation_crops_train[i]
            cumulative_price_decrease *= 1 - price_decrease_mushroom_train[i]
            yangdu *= 0.95
            cost *= 1.05
            vegetable_price *= 1.05
            year_data[year] = {
                "玉米销量": round(cumulative_growth_wheat_corn, 4),
                "其他销量": round(cumulative_sales_variation, 4),
                "亩产量": round(cumulative_yield_variation, 4),
                "食用菌价格": round(cumulative_price_decrease, 4),
                "蔬菜价格": round(vegetable_price, 4),
                "羊肚菌价格": round(yangdu, 4),
                "种植成本": round(cost, 4),
            }
        cumulative_samples_list_train.append(year_data)

    # 第二次采样 - 用于模型评价
    lhs_samples_eval = lhs(n_variables, samples=n_eval_samples)

    growth_rate_wheat_corn_eval = 0.05 + lhs_samples_eval[:, 0] * 0.05
    sales_variation_other_crops_eval = -0.05 + lhs_samples_eval[:, 1] * 0.10
    yield_variation_crops_eval = -0.10 + lhs_samples_eval[:, 2] * 0.20
    price_decrease_mushroom_eval = 0.01 + lhs_samples_eval[:, 3] * 0.04

    # 生成包含字典的列表，用于模型评价
    cumulative_samples_list_eval = []
    for i in range(n_eval_samples):
        year_data = {}
        cumulative_growth_wheat_corn = 1.0
        cumulative_sales_variation = 1.0
        cumulative_yield_variation = 1.0
        cumulative_price_decrease = 1.0
        yangdu = 1.0
        cost = 1.0
        vegetable_price = 1.0
        for year in years:
            cumulative_growth_wheat_corn *= 1 + growth_rate_wheat_corn_train[i]
            cumulative_sales_variation *= 1 + sales_variation_other_crops_train[i]
            cumulative_yield_variation *= 1 + yield_variation_crops_train[i]
            cumulative_price_decrease *= 1 - price_decrease_mushroom_train[i]
            yangdu *= 0.95
            cost *= 1.05
            vegetable_price *= 1.05
            year_data[year] = {
                "玉米销量": round(cumulative_growth_wheat_corn, 4),
                "其他销量": round(cumulative_sales_variation, 4),
                "亩产量": round(cumulative_yield_variation, 4),
                "食用菌价格": round(cumulative_price_decrease, 4),
                "蔬菜价格": round(vegetable_price, 4),
                "羊肚菌价格": round(yangdu, 4),
                "种植成本": round(cost, 4),
            }
        cumulative_samples_list_eval.append(year_data)

    # 初始化一个空的 DataFrame 存储结果，行表示样本，列表示模型
    results_df = pd.DataFrame(
        np.zeros((n_eval_samples, n_train_samples)),
        columns=[f"Model_{i+1}" for i in range(n_train_samples)],
        index=[f"Sample_{i+1}" for i in range(n_eval_samples)],
    )

    # 使用partial绑定函数参数
    from functools import partial

    train_and_evaluate_partial = partial(
        train_and_evaluate,
        cumulative_samples_list_train=cumulative_samples_list_train,
        cumulative_samples_list_eval=cumulative_samples_list_eval,
        n_eval_samples=n_eval_samples,
    )

    # 并行处理训练和评价
    with ProcessPoolExecutor() as executor:
        all_results = list(
            executor.map(train_and_evaluate_partial, range(n_train_samples))
        )

    # 填充结果到 DataFrame
    for i in range(n_train_samples):
        results_df.iloc[:, i] = all_results[i]

    # 输出结果
    print(results_df)
    results_df.to_csv("model_results.csv", float_format="%.2f")

    # 计算最大后悔值和平均期望值
    max_regret_values = []
    average_expected_values = []

    for i in range(n_train_samples):
        optimal_values = results_df.max(axis=1)  # 假设值越大越好
        regret_values = optimal_values - results_df.iloc[:, i]
        max_regret = regret_values.max()  # 最大后悔值
        average_expected_value = results_df.iloc[:, i].mean()  # 平均期望值

        max_regret_values.append(max_regret)
        average_expected_values.append(average_expected_value)

    # 输出所有模型的最大后悔值和平均期望值
    for i in range(n_train_samples):
        print(
            f"Model {i+1}: 最大后悔值 = {max_regret_values[i]:.4f}, 平均期望值 = {average_expected_values[i]:.4f}"
        )
        # 保存最大后悔值和平均期望值到 DataFrame
    summary_df = pd.DataFrame(
        {
            "Model": [f"Model_{i+1}" for i in range(n_train_samples)],
            "最大后悔值": max_regret_values,
            "平均期望值": average_expected_values,
        }
    )

    # 保存到 CSV 文件
    summary_df.to_csv("model_summary.csv", index=False, float_format="%.4f")
    end_time = time.time()
    print(time_start - end_time, "s")
